from neo4j import GraphDatabase
from api.config.database import Config
import json
import os
import logging

class Neo4jService:
    def __init__(self):
        self._driver = GraphDatabase.driver(
            Config.NEO4J_URI, 
            auth=(Config.NEO4J_USERNAME, Config.NEO4J_PASSWORD)
        )
        self._initialized = False
    
    def close(self):
        self._driver.close()
    
    def get_medical_knowledge(self, query):
        """
        从知识图谱中检索医学知识
        Retrieve medical knowledge from the knowledge graph
        """
        with self._driver.session() as session:
            result = session.read_transaction(self._search_knowledge, query)
            return result
    
    @staticmethod
    def _search_knowledge(tx, query):
        # 这里是一个示例Cypher查询，根据实际图结构进行调整
        cypher_query = """
        MATCH (n)
        WHERE n.name CONTAINS $query OR n.description CONTAINS $query
        RETURN n LIMIT 10
        """
        result = tx.run(cypher_query, parameters={"query": query})
        return [record["n"] for record in result]
    
    def save_feedback(self, question, answer, feedback_score):
        """
        保存用户反馈数据用于模型微调
        Save user feedback data for model fine-tuning
        """
        with self._driver.session() as session:
            session.write_transaction(self._store_feedback, question, answer, feedback_score)
    
    @staticmethod
    def _store_feedback(tx, question, answer, feedback_score):
        cypher_query = """
        CREATE (f:Feedback {
            question: $question,
            answer: $answer,
            score: $feedback_score,
            timestamp: timestamp()
        })
        RETURN f
        """
        result = tx.run(cypher_query, 
                       parameters={"question": question, 
                                "answer": answer, 
                                "feedback_score": feedback_score})
        return result.single()
    
    def initialize_database(self, medical_data_path=None, knowledge_graph_path=None):
        """
        初始化Neo4j数据库，导入医疗数据和知识图谱
        Initialize Neo4j database with medical data and knowledge graph
        """
        try:
            # 默认路径
            if medical_data_path is None:
                medical_data_path = os.path.join(os.path.dirname(__file__), '..', 'data', 'medical_data.json')
            if knowledge_graph_path is None:
                knowledge_graph_path = os.path.join(os.path.dirname(__file__), '..', 'data', 'knowledge_graph.json')
            
            logging.info("开始初始化Neo4j数据库...")
            
            # 清空现有数据（可选）
            self._clear_database()
            
            # 导入医疗数据
            self._import_medical_data(medical_data_path)
            
            # 导入知识图谱
            self._import_knowledge_graph(knowledge_graph_path)
            
            # 创建索引
            self._create_indexes()
            
            self._initialized = True
            logging.info("Neo4j数据库初始化完成")
            
        except Exception as e:
            logging.error(f"初始化Neo4j数据库失败: {str(e)}")
            raise
    
    def _clear_database(self):
        """
        清空数据库
        Clear the database
        """
        with self._driver.session() as session:
            session.write_transaction(self._clear_all_data)
    
    @staticmethod
    def _clear_all_data(tx):
        """清空所有数据"""
        tx.run("MATCH (n) DETACH DELETE n")
    
    def _import_medical_data(self, medical_data_path):
        """
        导入医疗数据
        Import medical data from JSON file
        """
        if not os.path.exists(medical_data_path):
            logging.warning(f"医疗数据文件不存在: {medical_data_path}")
            return
        
        with open(medical_data_path, 'r', encoding='utf-8') as f:
            medical_data = json.load(f)
        
        with self._driver.session() as session:
            # 导入疾病数据
            diseases = medical_data.get('diseases', [])
            for disease in diseases:
                session.write_transaction(self._create_disease, disease)
            
            # 导入科室数据
            departments = medical_data.get('departments', [])
            for department in departments:
                session.write_transaction(self._create_department, department)
        
        logging.info(f"已导入 {len(diseases)} 个疾病和 {len(departments)} 个科室")
    
    @staticmethod
    def _create_disease(tx, disease):
        """创建疾病节点"""
        # 创建疾病节点
        disease_query = """
        CREATE (d:Disease {
            name: $name,
            description: $description
        })
        RETURN d
        """
        tx.run(disease_query, 
               name=disease['name'], 
               description=disease.get('description', ''))
        
        # 创建症状节点和关系
        symptoms = disease.get('symptoms', [])
        for symptom in symptoms:
            symptom_query = """
            MERGE (s:Symptom {name: $symptom_name})
            WITH s
            MATCH (d:Disease {name: $disease_name})
            CREATE (d)-[:HAS_SYMPTOM]->(s)
            """
            tx.run(symptom_query, 
                   symptom_name=symptom, 
                   disease_name=disease['name'])
        
        # 创建药物节点和关系
        drugs = disease.get('drugs', [])
        for drug in drugs:
            drug_query = """
            MERGE (dr:Drug {
                name: $drug_name,
                usage: $usage,
                indication: $indication
            })
            WITH dr
            MATCH (d:Disease {name: $disease_name})
            CREATE (d)-[:TREATED_BY]->(dr)
            """
            tx.run(drug_query,
                   drug_name=drug['name'],
                   usage=drug.get('usage', ''),
                   indication=drug.get('indication', ''),
                   disease_name=disease['name'])
    
    @staticmethod
    def _create_department(tx, department):
        """创建科室节点"""
        # 创建科室节点
        dept_query = """
        CREATE (dep:Department {name: $name})
        RETURN dep
        """
        tx.run(dept_query, name=department['name'])
        
        # 建立科室与症状的关系
        symptoms = department.get('symptoms', [])
        for symptom in symptoms:
            symptom_dept_query = """
            MERGE (s:Symptom {name: $symptom_name})
            WITH s
            MATCH (dep:Department {name: $dept_name})
            CREATE (s)-[:RELATED_TO]->(dep)
            """
            tx.run(symptom_dept_query,
                   symptom_name=symptom,
                   dept_name=department['name'])
    
    def _import_knowledge_graph(self, knowledge_graph_path):
        """
        导入知识图谱数据
        Import knowledge graph from JSON file
        """
        if not os.path.exists(knowledge_graph_path):
            logging.warning(f"知识图谱文件不存在: {knowledge_graph_path}")
            return
        
        with open(knowledge_graph_path, 'r', encoding='utf-8') as f:
            kg_data = json.load(f)
        
        with self._driver.session() as session:
            # 导入节点
            nodes = kg_data.get('nodes', [])
            for node in nodes:
                session.write_transaction(self._create_kg_node, node)
            
            # 导入关系
            relationships = kg_data.get('relationships', [])
            for rel in relationships:
                session.write_transaction(self._create_kg_relationship, rel)
        
        logging.info(f"已导入知识图谱: {len(nodes)} 个节点, {len(relationships)} 个关系")
    
    @staticmethod
    def _create_kg_node(tx, node):
        """创建知识图谱节点"""
        node_type = node.get('type', 'Unknown')
        node_id = node.get('id')
        node_name = node.get('name')
        
        # 根据类型创建不同的标签
        if node_type == '疾病':
            label = 'Disease'
        elif node_type == '症状':
            label = 'Symptom'
        elif node_type == '药物':
            label = 'Drug'
        elif node_type == '科室':
            label = 'Department'
        else:
            label = 'Entity'
        
        query = f"""
        MERGE (n:{label} {{
            id: $id,
            name: $name,
            type: $type,
            description: $description,
            usage: $usage
        }})
        RETURN n
        """
        
        tx.run(query,
               id=node_id,
               name=node_name,
               type=node_type,
               description=node.get('description', ''),
               usage=node.get('usage', ''))
    
    @staticmethod
    def _create_kg_relationship(tx, relationship):
        """创建知识图谱关系"""
        source_id = relationship.get('source')
        target_id = relationship.get('target')
        rel_type = relationship.get('type', 'RELATED_TO')
        description = relationship.get('description', '')
        
        # 将关系类型转换为大写，符合Neo4j惯例
        rel_type = rel_type.upper()
        
        query = f"""
        MATCH (s {{id: $source_id}})
        MATCH (t {{id: $target_id}})
        CREATE (s)-[r:{rel_type} {{description: $description}}]->(t)
        RETURN r
        """
        
        tx.run(query,
               source_id=source_id,
               target_id=target_id,
               description=description)
    
    def _create_indexes(self):
        """
        创建索引以提高查询性能
        Create indexes for better query performance
        """
        with self._driver.session() as session:
            indexes = [
                "CREATE INDEX IF NOT EXISTS FOR (d:Disease) ON (d.name)",
                "CREATE INDEX IF NOT EXISTS FOR (s:Symptom) ON (s.name)",
                "CREATE INDEX IF NOT EXISTS FOR (dr:Drug) ON (dr.name)",
                "CREATE INDEX IF NOT EXISTS FOR (dep:Department) ON (dep.name)",
                "CREATE INDEX IF NOT EXISTS FOR (n) ON (n.id)"
            ]
            
            for index_query in indexes:
                try:
                    session.run(index_query)
                except Exception as e:
                    logging.warning(f"创建索引失败: {str(e)}")
    
    def check_database_status(self):
        """
        检查数据库状态
        Check database status
        """
        try:
            with self._driver.session() as session:
                # 检查节点数量
                result = session.run("""
                MATCH (n)
                RETURN 
                    count(n) as total_nodes,
                    count(CASE WHEN 'Disease' IN labels(n) THEN 1 END) as diseases,
                    count(CASE WHEN 'Symptom' IN labels(n) THEN 1 END) as symptoms,
                    count(CASE WHEN 'Drug' IN labels(n) THEN 1 END) as drugs,
                    count(CASE WHEN 'Department' IN labels(n) THEN 1 END) as departments
                """)
                
                record = result.single()
                if not record:
                    return {'error': '无法获取数据库统计信息'}
                
                # 检查关系数量
                rel_result = session.run("MATCH ()-[r]->() RETURN count(r) as total_relationships")
                rel_record = rel_result.single()
                rel_count = rel_record['total_relationships'] if rel_record else 0
                
                status = {
                    'initialized': self._initialized,
                    'total_nodes': record['total_nodes'],
                    'diseases': record['diseases'],
                    'symptoms': record['symptoms'],
                    'drugs': record['drugs'],
                    'departments': record['departments'],
                    'total_relationships': rel_count
                }
                
                return status
                
        except Exception as e:
            logging.error(f"检查数据库状态失败: {str(e)}")
            return {'error': str(e)}
    
    def is_initialized(self):
        """
        检查数据库是否已初始化
        Check if database is initialized
        """
        try:
            status = self.check_database_status()
            if isinstance(status, dict) and 'total_nodes' in status:
                total_nodes = status.get('total_nodes', 0)
                return isinstance(total_nodes, int) and total_nodes > 0
            return False
        except:
            return False
    
    def query_disease_symptoms(self, disease_name: str):
        """
        查询特定疾病的症状
        Query symptoms of a specific disease
        """
        with self._driver.session() as session:
            result = session.read_transaction(self._get_disease_symptoms, disease_name)
            return result
    
    @staticmethod
    def _get_disease_symptoms(tx, disease_name):
        cypher_query = """
        MATCH (d:Disease {name: $disease_name})-[:HAS_SYMPTOM]->(s:Symptom)
        RETURN d.name as disease, d.description as description, collect(s.name) as symptoms
        """
        result = tx.run(cypher_query, disease_name=disease_name)
        record = result.single()
        if record:
            return {
                'disease': record['disease'],
                'description': record['description'],
                'symptoms': record['symptoms']
            }
        return None
    
    def query_symptom_diseases(self, symptom_name: str):
        """
        查询可能导致特定症状的疾病
        Query diseases that may cause a specific symptom
        """
        with self._driver.session() as session:
            result = session.read_transaction(self._get_symptom_diseases, symptom_name)
            return result
    
    @staticmethod
    def _get_symptom_diseases(tx, symptom_name):
        cypher_query = """
        MATCH (d:Disease)-[:HAS_SYMPTOM]->(s:Symptom {name: $symptom_name})
        RETURN s.name as symptom, collect({name: d.name, description: d.description}) as diseases
        """
        result = tx.run(cypher_query, symptom_name=symptom_name)
        record = result.single()
        if record:
            return {
                'symptom': record['symptom'],
                'diseases': record['diseases']
            }
        return None
    
    def query_disease_treatment(self, disease_name: str):
        """
        查询疾病的治疗方法和药物
        Query treatment methods and drugs for a disease
        """
        with self._driver.session() as session:
            result = session.read_transaction(self._get_disease_treatment, disease_name)
            return result
    
    @staticmethod
    def _get_disease_treatment(tx, disease_name):
        cypher_query = """
        MATCH (d:Disease {name: $disease_name})-[:TREATED_BY]->(dr:Drug)
        RETURN d.name as disease, d.description as description,
               collect({name: dr.name, usage: dr.usage, indication: dr.indication}) as drugs
        """
        result = tx.run(cypher_query, disease_name=disease_name)
        record = result.single()
        if record:
            return {
                'disease': record['disease'],
                'description': record['description'],
                'drugs': record['drugs']
            }
        return None
    
    def query_department_recommendation(self, symptoms: list):
        """
        根据症状推荐科室
        Recommend departments based on symptoms
        """
        with self._driver.session() as session:
            result = session.read_transaction(self._get_department_recommendation, symptoms)
            return result
    
    @staticmethod
    def _get_department_recommendation(tx, symptoms):
        cypher_query = """
        MATCH (s:Symptom)-[:RELATED_TO]->(dep:Department)
        WHERE s.name IN $symptoms
        RETURN collect(DISTINCT dep.name) as departments,
               collect(DISTINCT s.name) as matched_symptoms
        """
        result = tx.run(cypher_query, symptoms=symptoms)
        record = result.single()
        if record:
            return {
                'departments': record['departments'],
                'matched_symptoms': record['matched_symptoms']
            }
        return None
    
    def query_drug_information(self, drug_name: str):
        """
        查询药物信息
        Query drug information
        """
        with self._driver.session() as session:
            result = session.read_transaction(self._get_drug_information, drug_name)
            return result
    
    @staticmethod
    def _get_drug_information(tx, drug_name):
        cypher_query = """
        MATCH (dr:Drug {name: $drug_name})<-[:TREATED_BY]-(d:Disease)
        RETURN dr.name as drug, dr.usage as usage, dr.indication as indication,
               collect(d.name) as treats_diseases
        """
        result = tx.run(cypher_query, drug_name=drug_name)
        record = result.single()
        if record:
            return {
                'drug': record['drug'],
                'usage': record['usage'],
                'indication': record['indication'],
                'treats_diseases': record['treats_diseases']
            }
        return None
    
    def query_complex_relationship(self, question_keywords: list):
        """
        复杂关系查询，支持多跳查询
        Complex relationship query supporting multi-hop queries
        """
        with self._driver.session() as session:
            result = session.read_transaction(self._get_complex_relationship, question_keywords)
            return result
    
    @staticmethod
    def _get_complex_relationship(tx, keywords):
        # 查找疾病-症状-科室的完整路径
        cypher_query = """
        MATCH path = (d:Disease)-[:HAS_SYMPTOM]->(s:Symptom)-[:RELATED_TO]->(dep:Department)
        WHERE d.name IN $keywords OR s.name IN $keywords OR dep.name IN $keywords
        RETURN d.name as disease, s.name as symptom, dep.name as department,
               d.description as disease_description
        LIMIT 20
        """
        result = tx.run(cypher_query, keywords=keywords)
        relationships = []
        for record in result:
            relationships.append({
                'disease': record['disease'],
                'symptom': record['symptom'],
                'department': record['department'],
                'disease_description': record['disease_description']
            })
        return relationships

# Global instance
neo4j_service = Neo4jService()